#!/bin/bash
function showhelp() {
echo 'usage: gencp2k.sh -t -m "-3" opt *.xsd/*.cif'
echo 'requies vaspkit and Multiwfn'
echo 'option:-[t task, default:3. 1:energy, 2:ene+force, 3:opt, 4:opt+cell, 5:freq, 6:MD(NVT), 7:TS, 9:NMR, 10:Polar]'
echo '       -[m method, default:2. 1:LDA, 2:PBE, -2:revPBE, 3:TPSS, -3:PBEsol, -6:PBE0, -7:B3LYP, -8:HSE06, 30:GFN1-xTB]'
echo '       -[b basis, default:2. 1:SZV-MOLOPT-SR-GTH, 2:DZVP-MOLOPT-SR-GTH, 10:6-31G*, 20:cc-DZ]'
echo '       -[s scf, default:4. 4:OT+outer_scf, 45:OT, 6:Diagonal+smear, 52:Diagonal+Broyden, 516:Diagonal+DIIS+smear]'
echo '       -[c charge and mult, default:"0 1".]'
echo '       -[S sccs solvent mode, default:no. set epsilon of solvent]'
echo '       -[d surface dipole, default:0. 0:not use, 1:x, 2:y, 3:z]'
echo '       -[u dft+u, default off, use -u to turn it on ]'
echo '       -[P periodicity, default:XYZ. X, XY, XYZ, XZ, Y, YZ, Z]' 
echo '       -[k K-points, default: 1,1,1(gamma)] '
echo '       -[K temperateu for MD, default: 298.15] '
echo '       -[f fix atom, default: no atom or read fix atom infor from POSCAR/.xsd file]'
echo '         format: 1,3-15 for atom index or 0.2 to fix atom with z fraction coord < 0.2'
echo '         or "no" to force not fix atom even if POSCAR/.xsd file has fixed atom info'
echo '       -[v number of virtual orbital for dos, default: None, to set this number scf should not be OT]'
echo '       -[C calc charge, default:0. 0:none, 1:Mulliken, 2:Lowdin, 3:Hirshfeld, 4:Hirshfeld-i, 5:voronoi, 6:resp, 7:REPEAT]'
echo '       -[M generate molden for multiwfn]'
echo '       -[G out cube, default:0. 1:density, 2:ELF, 3:XCpot, 4:Hatree_pot, 5:E-field, 6:MO, 7:1+4]' 
echo '       -[T use Multiwfn input template, you could use MULTIWFN_INPUT file from previous project]' 
echo '       -[O set max cyc for innter and outer scf cycle]' 
exit
}
if [ $# -lt 1 ]; then
    showhelp
    fi
while getopts ":t:m:b:s:c:v:S:d:P:k:K:f:C:MG:T:u:O:" opt
do
    case $opt in
        t) TASK=$OPTARG;;
        m) METHOD=$OPTARG;;
        b) BASIS=$OPTARG;;
        s) SCF=$OPTARG;;
        c) CHARGEMULT=$OPTARG;;
        S) SCCS=$OPTARG;;
        v) VIRTUAL=$OPTARG;;
        d) SURF_DIPOLE=$OPTARG;;
        P) PERIOD=$OPTARG;;
        k) KPOINT=$OPTARG;;
        K) TEMPERATURE=$OPTARG;;
        f) FIXATOM=$OPTARG;;
        C) CHARGETYPE=$OPTARG;;
        M) MOLDEN=on;;
        G) CUBETYPE=$OPTARG;;
        T) TEMPLATE=$OPTARG;;
        O) OTCYC=$OPTARG;;
        u) DFTU=on;;
        ?) showhelp
           exit 1;;
    esac
done
shift $(($OPTIND - 1))

dos2unix $1

# set default parameter
if [[ -z "$TASK" ]];then TASK=3;fi
if [[ -z "$METHOD" ]];then METHOD=2;fi
if [[ -z "$BASIS" ]];then BASIS=2;fi
if [[ -z "$VDW" ]];then VDW=2;fi
if [[ -z "$SCF" ]];then 
    if [[ "$METHOD" == "30" ]];then
	    SCF="5"
    else
	    SCF="4"
    fi
fi
if [[ -z "$CHARGEMULT" ]];then CHARGEMULT="0 1";fi
if [[ -z "$KPOINT" ]];then KPOINT=1,1,1;fi
if [[ -z "$STRUCFORMAT" ]];then STRUCFORMAT=3;fi
if [[ -z "$OTCYC" ]];then OTCYC="90 3";fi

# parser some parameter
CHARGE=`echo $CHARGEMULT | cut -d " " -f 1`
MULTI=`echo $CHARGEMULT | cut -d " " -f 2`
INNER_CYC=`echo $OTCYC | cut -d " " -f 1`
OUTER_CYC=`echo $OTCYC | cut -d " " -f 2`
INPNAME=INPNAME
function gen_minp {
    if [ -n "$TEMPLATE" ];then
        if [ -f "../$TEMPLATE" ];then
            cp ../$TEMPLATE ./MULTIWFN_INPUT
            sed -i "2d" MULTIWFN_INPUT
            sed -i "1a $INPNAME" MULTIWFN_INPUT
        else
            echo "Error! template file $TEMPLATE not found"
            exit
        fi
    else
        echo "cp2k" > MULTIWFN_INPUT
        echo "$INPNAME" >> MULTIWFN_INPUT
        echo -en "-9\n1\n" >> MULTIWFN_INPUT
        echo "$CHARGE" >> MULTIWFN_INPUT
        # set multiplicity
        echo -en "2\n$MULTI\n" >> MULTIWFN_INPUT
        # print level medium
        echo -en "10\n2\n" >> MULTIWFN_INPUT
        if [[ -n "$SURF_DIPOLE" ]];then
            echo "23" >> MULTIWFN_INPUT
            echo "$SURF_DIPOLE" >> MULTIWFN_INPUT
        fi
        # show orbital level
        echo "14" >> MULTIWFN_INPUT
        if [[ "$DFTU" == "on" ]];then echo 8 >> MULTIWFN_INPUT;fi
        if [[ -n "$VIRTUAL" ]];then
          echo -en "12\n$VIRTUAL\n" >> MULTIWFN_INPUT
        fi
        echo "0" >> MULTIWFN_INPUT
        # Set direction(s) of applying periodic boundary condition
        if [[ -n "$PERIOD" ]] && [[ "PERIOD" != "XYZ" ]];then
            echo -en "-7\n$PERIOD\n" >> MULTIWFN_INPUT
        fi
        # Set charge type to calculate 
        if [[ -n "$CHARGETYPE" ]];then
            echo -en "-4\n$CHARGETYPE\n" >> MULTIWFN_INPUT
        fi
        # Set exporting cube file 
        if [[ -n "$CUBETYPE" ]];then
            echo -en "-3\n$CUBETYPE\n" >> MULTIWFN_INPUT
        fi
        # default exporting molden file
        # echo "-2" >> MULTIWFN_INPUT
        # choose task
        echo "-1" >> MULTIWFN_INPUT
        echo "$TASK" >> MULTIWFN_INPUT
        # v-rescale thermostat (NVT)
        if [[ "$TASK" == "6" ]];then
            echo -en "10\n2\n" >> MULTIWFN_INPUT
        fi
        if [[ "$MOLDEN" == "on" ]];then
          echo -2 >> MULTIWFN_INPUT
        fi
        # choose theoretical method
        echo "1" >> MULTIWFN_INPUT
        echo "$METHOD" >> MULTIWFN_INPUT
        # choose basis set
        echo -en "2\n$BASIS\n" >> MULTIWFN_INPUT
        # choose dispersion correction
        echo -en "3\n$VDW\n" >> MULTIWFN_INPUT
        # chooe SCF strategy
        echo -e "$SCF" >> MULTIWFN_INPUT
        # SCCS
        if [[ -n "$SCCS" ]];then
          echo "7" >> MULTIWFN_INPUT
        fi
        # KPOINT 
        if [[ -n "$KPOINT" ]];then
            echo -en "8\n$KPOINT\n" >> MULTIWFN_INPUT
        fi
        # outputted structure format
        echo "-5" >> MULTIWFN_INPUT
        echo "$STRUCFORMAT" >> MULTIWFN_INPUT
        if [[ -n "$FIXATOM_IDX" ]];then
           echo "9" >> MULTIWFN_INPUT
           echo "$FIXATOM_IDX" >> MULTIWFN_INPUT
        fi  
        echo "0" >> MULTIWFN_INPUT
        echo "q" >> MULTIWFN_INPUT
    fi
}
declare -A task2str=( ["1"]="SP" ["2"]="FORCE" ["3"]="OPT" ["4"]="OPTC" ["5"]="FREQ" ["6"]="MD" ["7"]="TS" ["9"]="NMR" ["10"]="POLAR")

function gen_fixatom {
    if [[ $FIXATOM == "no" ]];then 
        FIXATOM_IDX=""
        return 0
    fi
    if [[ $FIXATOM == *"\<"* || $FIXATOM == *">"* ]];then
        uth=${FIXATOM#*<}     # 删除"<"前的所有字符
        uth=${uth%%[^0-9.]*}   # 删除浮点数后的所有字符
        lth=${FIXATOM#*>}     # 删除">"前的所有字符
        lth=${lth%%[^0-9.]*}   # 删除浮点数后的所有字符
        echo -en "100\n2\n33\n\n0\nq" | Multiwfn $1 &> /dev/null
        flnm=${1%.*}.cif
        FIXATOM_IDX=`awk -v lth=$lth -v uth=$uth 'BEGIN{idx=0;f=""}{if(NF==4){idx++;if($4<uth && $4>lth){f=f""idx","}}}END{print substr(f,1,length(f)-1)}' $flnm`
    elif [ -f "POSCAR" ];then
        echo "$1":
        a1=(`sed -n '6p' POSCAR`)
        a2=(`sed -n '7p' POSCAR`)
        echo -n "Formula: "
        for ((i=0;i<${#a1[@]};i++)); do echo -n ${a1[$i]}:${a2[$i]}" "; done
        echo ""
        echo "Charge: $CHARGE    Multiplicity: $MULTI"
        FIXATOM_IDX=`awk '{if(NR==9 && $1 ~/^[DdCc]/){C=1;next}}{if(C==1)fix[$4$5$6]=fix[$4$5$6]","NR-9}END{print substr(fix["FFF"],2,length(fix["FFF"]))}' POSCAR`
    fi
    nfixatom=`echo $FIXATOM_IDX | sed 's/,/ /g' | wc -w`
    if [[ -n "$nfixatom" ]];then
        echo "Fix $nfixatom atoms"
    fi
    echo $FIXATOM_IDX
}


while [[ $# -ge 1 ]];do
    basename=${1%.*}_${task2str[$TASK]}
    INPNAME=${basename}.inp
    mkdir ${basename} 2> /dev/null
    cp $1 ${basename}
    cd ${basename}
    PWD=`pwd`
    suffix=${1##*.}
    if [[ $suffix == "xsd" ]];then
        echo -e "106\n$1" | vaspkit 1> /dev/null 
        inp=POSCAR
    elif [[ $suffix == "pdb" || $suffix == "cif" || $suffix == "gro" ]];then
        inp=$1
    else
        echo "unrecognized file $1, pass"
        continue
    fi
    gen_fixatom $inp
    gen_minp
    cat MULTIWFN_INPUT | Multiwfn $inp 1> /dev/null
    sed -i "/steps of inner SCF/c\      MAX_SCF $INNER_CYC" $INPNAME
    sed -i "/steps of outer SCF/c\        MAX_SCF $OUTER_CYC" $INPNAME
    sed -i '/STEPS 2/c\    STEPS 5000' $INPNAME
    if [ -n "$TEMPERATURE" ];then
        sed -i "/TEMPERATURE/c\    TEMPERATURE $TEMPERATURE" $INPNAME
    fi
    echo "$INPNAME generated in $basename"
    cd ..
    shift 1
done

